import torch
import numpy as np
import pandas as pd
import torch.nn as nn #导入nn，nn.linear是一个全连接层
from transformers import BertModel, BertTokenizer# 导入bert模型和bert的分词器
from sklearn.model_selection import train_test_split #sklearn是一个非常基础的机器学习库，里面包括聚类算法、逻辑回归算法等，train_test_split方法是把数据一部分用作训练一部分用作测试的划分数据的方法

# 第一部分：加载训练集
# 第一个参数说明文件位置，默认会以空格划分每行的内容，
# 第二个delimier参数设置的备选用制表符划分，
# 第三个参数表示不满足要求的行舍弃
train_set = pd.read_csv("D:\study\毕设\data\question_train_dataset.csv",error_bad_lines=False)
test_set= pd.read_csv("D:\study\毕设\data\question_test_dataset.csv",error_bad_lines=False)
# 下面是bert模型中的hfl/chinese-bert-wwm模型的模型以及模型分词方法
# model_name = "hfl/chinese-bert-wwm"
# tokenizer = BertTokenizer.from_pretrained(model_name)
# model = BertModel.from_pretrained(model_name)

# 第二部分
# 采用bert微调策略，在反向传播时一同调整BERT和线性层的参数，使bert更加适合分类任务
# bert微调，在这里 就是bert（类的前四行）+一个全连接层（第五行 self.fc = nn.Linear(768,15)）组成一个全新模型，
class BertClassfication(nn.Module):
    def __init__(self):
        # 前四行就是bert，第五行是全连接层，即bert微调
        super(BertClassfication,self).__init__()
        self.model_name = 'hfl/chinese-bert-wwm'
        self.model = BertModel.from_pretrained(self.model_name)
        self.tokenizer = BertTokenizer.from_pretrained(self.model_name)
        self.fc = nn.Linear(768,15)     #768取决于BERT结构，12-layer, 768-hidden, 12-heads, 110M parameters

    def forward(self,x):#这里的输入x是一个list,也就是输入文本x
        # 这句话是对文本1.进行分词，2.进行编码 详见https://www.cnblogs.com/zjuhaohaoxuexi/p/15135466.html
        # 第一个参数x表示文本
        # 第二个参数add_special_token表示要不要加入一些特殊字符，比如标记开头的cls等
        # 第三个参数表示最长长度
        # 第四个参数pad_to_max_len表示如果不够max_len的长度，就补长到最大长度
        # 还有一个没用到的参数参数：比最长的长久减掉，第四第五个参数也就是长的减，短的补
        batch_tokenized = self.tokenizer.batch_encode_plus(x, add_special_tokens=True,
                                max_length=148, pad_to_max_length=True)      #tokenize、add special token、pad

        # 取出input_ids：这是对汉语的编码
        # attention_mask:这是对每个字是否mask的一个标记，原本的词的位置是1，如果由于词长不够max_len，用pad或者其他的填充了，该位置就是0，表示在模型中不注意他，因为该位置是填充上的没有意义的字符
        input_ids = torch.tensor(batch_tokenized['input_ids'])
        attention_mask = torch.tensor(batch_tokenized['attention_mask'])
        # 把上边两个输入bert模型，得到bert最后一层的隐藏层的输出，
        hiden_outputs = self.model(input_ids,attention_mask=attention_mask)
        # bert的输出结果有四个维度：
        # last_hidden_state：shape是(batch_size, sequence_length, hidden_size)，hidden_size=768,它是模型最后一层输出的隐藏状态。
        # pooler_output：shape是(batch_size, hidden_size)，这是序列的第一个token(classification token)的最后一层的隐藏状态，它是由线性层和Tanh激活函数进一步处理的。（通常用于句子分类，至于是使用这个表示，还是使用整个输入序列的隐藏状态序列的平均化或池化，视情况而定）
        # hidden_states：这是输出的一个可选项，如果输出，需要指定config.output_hidden_states=True,它也是一个元组，它的第一个元素是embedding，其余元素是各层的输出，每个元素的形状是(batch_size, sequence_length, hidden_size)
        # attentions：这也是输出的一个可选项，如果输出，需要指定config.output_attentions=True,它也是一个元组，它的元素是每一层的注意力权重，用于计算self-attention heads的加权平均值。
        # 我们是微调模式，需要获取bert最后一个隐藏层的输出输入到下一个全连接层，所以取第一个维度，也就是hiden_outputs[0]
        # 此时shape是(batch_size, sequence_length, hidden_size)，[:,0,:]的意思是取出第一个也就是cls对应的结果，至于为什么这样操作，我也不知道，有人告诉我因为它最具有代表性，但为什么我还是不知道，有大神能给我讲一下吗
        outputs = hiden_outputs[0][:,0,:]     #[0]表示输出结果部分，[:,0,:]表示[CLS]对应的结果
        # 把bert最后一层的结果输入到全连接层中，全连接层是{768,15},会输出15分类的一个结果，
        output = self.fc(outputs)
        # 这里就是返回最终地分类结果了
        return output

# 第三部分，整理数据集
train_features = train_set['question'].values
train_targets = train_set['label'].values
test_features = test_set['question'].values
test_targets = test_set['label'].values
targets=np.hstack((train_targets, test_targets))
# train_features,test_features,train_targets,test_targets = train_test_split(sentences,targets)# 把数据分为训练集和测试集
batch_size = 32 # 把32个问题当成一个块来处理，相当于一段有32句话
batch_count = int(len(train_features) / batch_size) # 这里就是所有的数据一共可以分为多少块（段）
batch_train_inputs, batch_train_targets = [], [] # 一个列表存储分段的文本，一个列表存储分段的标签
# 分段
for i in range(batch_count):
    batch_train_inputs.append(train_features[i*batch_size : (i+1)*batch_size])
    batch_train_targets.append(train_targets[i*batch_size : (i+1)*batch_size])

# 第四部分，训练
bertclassfication = BertClassfication() #实例化
lossfuction = nn.CrossEntropyLoss() #定义损失函数，交叉熵损失函数
optimizer = torch.optim.Adam(bertclassfication.parameters(),lr=2e-5) #torch.optim里面包含一些优化器（一些反向传播调整参数的方法，比如梯度下降，随机梯度下降等），ADAM是一种随机梯度下降的改进优化方法
epoch = 5 # 训练轮数，即所有数据跑5遍
for _ in range(epoch):
    los = 0  # 损失，每一轮训练的损失应该重新开始计数
    for i in range(batch_count):# batch_count表示有多少块（段），每段有32句话
        inputs = batch_train_inputs[i]
        targets = torch.tensor(batch_train_targets[i])
        optimizer.zero_grad() # 1.梯度置零
        outputs= bertclassfication(inputs) # 2.模型获得结果
        loss = lossfuction(outputs,targets) # 3.计算损失
        loss.backward() # 4.反向传播
        optimizer.step() # 5.修改参数，w，b
        los += loss.item() #item()返回loss的值
        # 查看每处理五个段的损失
        if i%5==0:
            print("Batch:%d,Loss %.4f" % ((i),los/5))
            los = 5

# 第四部分：验证分类准确率
hit = 0 # 记录预测正确的个数
total = len(test_features) # 记录例子总数
for i in range(total):
    outputs = bertclassfication([test_features[i]])
    _,predict = torch.max(outputs,1)# torch.max函数详见https://www.jianshu.com/p/3ed11362b54f
    if predict==test_targets[i]:# 表示预测正确
        hit+=1
acc=round(hit/len(test_features),2)
print('准确率为%.4f'%(acc))

torch.save(bertclassfication,r'..\..\model\bertclassfication_'+str(acc)+'.pt')

# 模型训练完之后，由于原数据集的label只有数字，没有中文注释，人工注释0-14对应的类别
bertclassfication=torch.load(r'..\..\model\bertclassfication_0.93.pt')
transform_dict = {0:'杂豆病害',1:'杂豆虫害',2:'英文名',3:'别名',4:'简介',5:'为害症状',6:'病原物',
                  7:'侵染循环',8:'发生因素',9:'形态特征',10:'生活习性',11:'防治方法',12:'患病害杂豆',13:'患虫害杂豆',14:'致病病害'}
question=input("请输入问题：")
result = bertclassfication([question])
_,result = torch.max(result,1)
print(transform_dict[result.tolist()[0]])
